function R = dpmm_gibbs(base_gm, X, lalpha, N, varargin)
% Performs Gibbs sampling of a DP mixture model
%
%   R = dpmm_gibbs(base_gm, X, lalpha, N, ...);
%       performs Gibbs sampling of a (generalized) DP mixture model.
%
%       Inputs:
%       - base_gm:  the base generative model
%                   (which integrates both the base distribution of
%                    parameters as prior and the likelihood model)
%       - X:        the observed samples
%       - lalpha:   the log-weight of sampling from the base
%                   can be a scalar, or a pair [lalpha0, lalpha1],
%                   then lalpha0 is used in booting, and lalpha1 used in
%                   updating.
%       - N:        the number of samples to collect.
%
%       The output R is a struct with the following fields:
%       - tag:  a string 'mm_samples' (indicating what the struct is for)
%       - alg:  a string 'fmm_gibbs' (indicating how it is derived)
%       - n:    the number of samples
%       - samples:  a struct array, of which each element 
%                   corresponds to a sample and has the following fields:
%                   - Id:       row vector, which is 1:K
%                   - Theta:    model parameters [param.dim x K]
%                   - Pi:       the prior distribution [1 x K]
%                   - W:        the posterior weights of each component
%
%       In addition, the caller can specify the following
%       options in name/value list:
%
%       - 'BurnIn':     the number of iterations to burn in 
%                       default = 200.
%
%       - 'Interval':   the number of iterations to take before collecting
%                       next sample. default = 20
%
%       - 'Inherits':   the multi-prior object due to inheritance.
%                       (default = [], indicating no inheritance)
%
%       - 'PriW':       the prior weights of the inherited particles 
%
%       - 'PriQ':       the inheritance probabilities (default = [])
%
%       - 'Kmax':       the maximum allowable number of components
%                       default = inf;
%
%       - 'Display':    the level of displaying
%                       can be either of 'off', 'sample', 'move'
%

%   History
%   -------
%       - Created by Dahua Lin, on Nov 15, 2010


%% verify input arguments

if ~isobject(base_gm)
    error('dpmm_gibbs:invalidarg', ...
        'base_gm should be an object to represent a generative model.');
end

nx = base_gm.check_observations(X);
if nx < 0
    error('dpmm_gibbs:invalidarg', 'The observations X is not valid.');
end

if ~(isfloat(lalpha) && (isscalar(lalpha) || numel(lalpha) == 2) && isreal(lalpha))
    error('dpmm_gibbs:invalidarg', 'lalpha should be a real scalar or pair.');
end
lalpha = double(lalpha);

if isscalar(lalpha)
    lalpha0 = lalpha;
    lalpha1 = lalpha;
else
    lalpha0 = lalpha(1);
    lalpha1 = lalpha(2);
end


if ~(isnumeric(N) && isscalar(N) && N == fix(N) && N > 0)
    error('dpmm_gibbs:invalidarg', 'N should be a positive integer scalar.');
end

% check options

opts = struct( ...
    'BurnIn', 200, ...
    'Interval', 20, ...
    'Inherits', [], ...
    'PriW', zeros(1, 0), ...
    'PriQ', zeros(1, 0), ...
    'Kmax', inf, ...
    'Display', 'off');

if ~isempty(varargin)
    opts = parlist(opts, varargin{:});
    
    opts.BurnIn = check_pos_int('BurnIn', opts.BurnIn);
    opts.Interval = check_pos_int('Interval', opts.Interval);
    opts.Kmax = check_pos_int('Kmax', opts.Kmax);            
end

m0 = 0;
if ~isempty(opts.Inherits)
    m0 = opts.Inherits.num_priors;    
    if ~(isfloat(opts.PriW) && isvector(opts.PriW) && numel(opts.PriW) == m0)
        error('dpmm_gibbs:invalidarg', ...
            'PriW should be a numeric vector of length m0');
    end
    if ~(isfloat(opts.PriQ) && isvector(opts.PriQ) && numel(opts.PriQ) == m0)
        error('dpmm_gibbs:invalidarg', ...
            'PriQ should be a numeric vector of length m0');
    end
    
    priw = double(opts.PriW);
    if size(priw, 1) > 1; priw = priw.'; end
    
    priq = double(opts.PriQ);
    if size(priq, 1) > 1; priq = priq.'; end
else
    priw = zeros(1, 0);
    priq = zeros(1, 0);
end


switch lower(opts.Display)
    case 'sample'
        dplevel = 1;
    case 'move'
        dplevel = 2;
    case 'off'
        dplevel = 0;
    otherwise
        error('dpmm_gibbs:invalidarg', 'The Display option is invalid.');
end
        

%% main

% initialize

sys.base_gm = base_gm;
sys.inherits = opts.Inherits;
sys.m0 = m0;
sys.nx = nx;
sys.X = X;
sys.logp0 = base_gm.logmargin(X);
assert(isequal(size(sys.logp0), [1 nx]));
sys.priw = priw;
sys.priq = priq;

status = dpmm_init(sys, lalpha0, opts.Kmax);

if dplevel >= 1
    fprintf('Initialized => K = %d\n', numel(status.id));
end


% burn in

for t = 1 : opts.BurnIn    
    status = dpmm_move(sys, status, lalpha1, opts.Kmax);
    if dplevel >= 2
        disp_move(t, opts.BurnIn, status);
    end
end

if dplevel >= 1
    fprintf('Burned in => K = %d\n', numel(status.id));
end


% collect

samples = cell(N, 1);
samples{1} = make_sample(status);

if dplevel >= 1
    disp_sample(1, samples{1});
end


for i = 2 : N
    
    for t = 1 : opts.Interval
        status = dpmm_move(sys, status, lalpha1, opts.Kmax);
        if dplevel >= 2
            disp_move(t, opts.Interval, status);
        end
    end
    
    samples{i} = make_sample(status);   
    if dplevel >= 1
        disp_sample(i, samples{i});
    end
end

% output

R.tag = 'mm_samples';
R.alg = 'dpmm_gibbs';
R.n = N;
R.samples = vertcat(samples{:});


%% Core sampling functions

function s = dpmm_init(sys, lalpha, Kmax)
% status s is a struct with the following fields:
%
%   - z:  the assignment vector [1 x ns int32]
%   - id: the identifiers [1 x m int32], globally unique ID
%   - u:  the sources [1 x m int32], where the sample is obtained from
%   - c:  the counts [1 x m int32],
%   - w:  the weights [1 x m double]
%   - theta: the obtained particles [1 x m double]
%   - lik: the log-likelihood [ns x m double]
%   - max_id
%   
%   Here, m is the number of particles
%
   
nx = sys.nx;
s.z = zeros(1, nx, 'int32');

m0 = sys.m0;
pd = sys.base_gm.pdim;

s.id = int32(1:m0);
s.u = int32(1:m0);
s.c = zeros(1, m0, 'int32');
s.theta = zeros(pd, 0);
if m0 == 0
    s.w = zeros(1, 0);    
    s.lik = zeros(0, nx);    
else
    s.w = sys.priw .* sys.priq;   
    s.lik = sys.inherits.logmargin(sys.X);
    assert(isequal(size(s.lik), [m0 nx]));
end
s.max_id = int32(m0);

% first sequential sampling pass
s = seq_pass(sys, s, lalpha, Kmax);


function s = dpmm_move(sys, s, lalpha, Kmax)

m = numel(s.c);
base_gm = sys.base_gm;
inherits = sys.inherits;
X = sys.X;

% parameter update
grp = intgroup(m, s.z);
for k = 1 : m
    gk = grp{k};
    if ~isempty(gk)
        sX = base_gm.select_observations(X, gk);   
        uk = s.u(k);
        if uk == 0
            s.theta(:,k) = base_gm.pos_sample(sX, 1, 1);
        else
            s.theta(:,k) = inherits.pos_sample(sX, 1, 1, k);
        end          
    end
end

% label update
s = seq_pass(sys, s, lalpha, Kmax);


function s = seq_pass(sys, s, lalpha, Kmax)

nx = sys.nx;
m0 = numel(sys.priw);

ord = int32(randperm(nx));
rvs = rand(1, nx);

% do sampling

while ~isempty(ord)
    
    if numel(s.id) >= Kmax
        lalpha = -inf;
    end
    
    [z, c, w, i, k] = dpmm_seq_draw(sys, s, lalpha, ord, rvs);
    s.z = z;
    s.c = c;
    s.w = w;
    
    if ~isempty(k)  
        
        assert(k >= 0 && k <= m0);
        
        k0 = s.z(ord(i));
        if k0 > 0
            s.c(k0) = s.c(k0) - 1;
            s.w(k0) = s.w(k0) - 1;
            if s.c(k0) == 0
                if k0 <= m0
                    s.w(k0) = sys.priw(k0) * sys.q(k0);
                else
                    s.w(k0) = 0;
                end
            end                
        end

        cx = sys.base_gm.select_observations(sys.X, ord(i));
        if k == 0
            gm = sys.base_gm;            
            new_theta = gm.pos_sample(cx, 1, 1);
            new_lik = gm.loglik(new_theta, sys.X);
            
            s.max_id = s.max_id + 1;
            s.id = [s.id, s.max_id];
            s.u = [s.u, int32(0)];
            s.c = [s.c, int32(1)];
            s.w = [s.w, 1.0];
            s.theta = [s.theta, new_theta];
            s.lik = [s.lik; new_lik];
            s.z(ord(i)) = numel(s.id);
            
        else
            gm = sys.inherits;
            new_theta = gm.pos_sample(cx, 1, 1, k);
            new_lik = gm.loglik(new_theta, sys.X);
            
            s.c(k) = 1;
            s.w(k) = sys.priw(k) + 1;
            s.theta(:,k) = new_theta;
            s.lik(k, :) = new_lik;
            s.z(ord(i)) = k;
        end
    end
    
    ord = ord(i+1:end);
    rvs = rvs(i+1:end);
end

% filter
m = numel(s.id);
idx = 1:m;
can_rm = idx > m0 & s.w == 0;
if any(can_rm)
    rt = find(~can_rm);
    
    zmap = zeros(1, m, 'int32');
    zmap(rt) = 1 : numel(rt);
    
    s.z = zmap(s.z);
    
    s.id = s.id(rt);
    s.u = s.u(rt);
    s.c = s.c(rt);
    s.w = s.w(rt);
    s.theta = s.theta(:, rt);
    s.lik = s.lik(rt, :);
end



%% Auxiliary functions

function s = make_sample(status)

s = struct( ...
    'Id', status.id, ...
    'Theta', status.theta, ...
    'C', status.c, ...
    'W', status.w);


function v = check_pos_int(name, v)

if ~(isnumeric(v) && isscalar(v) && v > 0 && v == fix(v))
    error('dpmm_gibbs:invalidarg', ...
        '%s should be a positive integer scalar.', name);
end
v = double(v);


function disp_sample(i, s)

fprintf('sample %d: K = %d, max_id = %d\n', i, numel(s.Id), max(s.Id));


function disp_move(t, n, status)

fprintf('\tmove %d/%d: K = %d, max_id = %d\n', t, n, numel(status.id), status.max_id);



